# -*- coding: utf-8 -*-
"""
@author: Trae
@contact: traetai@gmail.com
@software: TraeAI
@file: detector.py
@time: 2024/7/29 16:30
@desc: 封装YOLOv8的检测器
"""
import torch
from ultralytics import YOLO

class Detector:
    """
    封装了YOLOv8模型的检测器类。
    """
    def __init__(self, model_path: str, device: str = 'auto'):
        """
        初始化检测器
        Args:
            model_path (str): YOLOv8模型的路径 (.pt文件)。
            device (str): 推理设备 ('cpu', 'cuda', or 'auto')。
        """
        self.model_path = model_path
        self.device = self._select_device(device)
        self.model = self._load_model()

    def _select_device(self, device: str) -> str:
        """根据用户选择和系统能力决定最终的设备。"""
        if device == 'auto':
            return 'cuda' if torch.cuda.is_available() else 'cpu'
        elif device == 'cuda' and not torch.cuda.is_available():
            print("Warning: CUDA not available, falling back to CPU.")
            return 'cpu'
        return device

    def _load_model(self) -> YOLO:
        """加载YOLOv8模型到指定设备。"""
        print(f"Loading model from {self.model_path} to {self.device}...")
        model = YOLO(self.model_path)
        model.to(self.device)
        print("Model loaded successfully.")
        return model

    def detect(self, frame, confidence_threshold: float = 0.4, target_classes: list = None):
        """
        对单帧图像进行对象检测。

        Args:
            frame: 输入的图像帧 (numpy array)。
            confidence_threshold (float): 置信度阈值。
            target_classes (list, optional): 需要检测的目标类别ID列表。如果为None，则检测所有类别。

        Returns:
            A list of detections from the model.
        """
        results = self.model(frame, conf=confidence_threshold, classes=target_classes, verbose=False)
        return results[0] # 返回第一个结果对象

if __name__ == '__main__':
    # 这是一个测试该模块功能的示例
    # 请确保你的项目结构中存在yolov8s.pt和一张测试图片
    import cv2
    from pathlib import Path

    # 假设项目根目录下有yolov8s/yolov8s.pt和video/test.jpg
    # 注意：这里的路径是相对于文件本身的，实际使用时应由主程序传入配置
    root_dir = Path(__file__).resolve().parents[2]
    model_path = root_dir / 'yolov8s' / 'yolov8s.pt'
    image_path = root_dir / 'video' / 'test.jpg' # 你需要一张名为test.jpg的测试图片

    if not model_path.exists() or not image_path.exists():
        print(f"Error: Ensure model '{model_path}' and image '{image_path}' exist.")
    else:
        # 1. 初始化检测器
        detector = Detector(str(model_path))

        # 2. 读取测试图片
        img = cv2.imread(str(image_path))

        # 3. 进行检测
        detections = detector.detect(img, confidence_threshold=0.3, target_classes=[0,2]) # 只检测人

        # 4. 打印并可视化结果
        print(f"Detected {len(detections.boxes)} objects.")
        annotated_frame = detections.plot() # yolov8自带的绘图功能

        cv2.imshow("Detections", annotated_frame)
        cv2.waitKey(0)
        cv2.destroyAllWindows()